"""Exact Riemann solver for the Euler Equations"""

from math import sqrt
import numpy

from euler import ideal_gas_w2q, ideal_gas_q2w, ideal_gas_flux_q

class RiemannSolverEulerExact(object):
    """Exact Riemann solver for the Euler equations.

    The Riemann solver used is defined in 'Riemann solvers and
    Numerical Methods for Fluid Dynamics' by Toro, Chapter 4

    """

    def __init__(self, gamma = 1.4, niter = 20, tol = 1e-6):
        """Constructor.

        Parameters:
        -----------

        gamma : float
            Ratio of specific heats.

        niter : int
            Maximum number of iterations for the Newton Raphson iteration

        tol : float
            Tolerance for the Newton Raphson iteration.

        """

        self.gamma = gamma

        # save derived variables
        tmp1 = 1.0/(2*gamma)
        tmp2 = 1.0/(gamma - 1.0)
        tmp3 = 1.0/(gamma + 1.0)

        self.gamma1 = (gamma - 1.0) * tmp1
        self.gamma2 = (gamma + 1.0) * tmp1
        self.gamma3 = 2*gamma * tmp2
        self.gamma4 = 2 * tmp2
        self.gamma5 = 2 * tmp3
        self.gamma6 = tmp3/tmp2
        self.gamma7 = 0.5 * (gamma - 1.0)
        self.gamma8 = gamma - 1.0

        self.tol = tol
        self.niter = niter

        # number of waves in the Riemann solver solution 
        self.nwaves = 3

    def compute_star(self, rhol, rhor, pl, pr, ul, ur):
        """Solve the Riemann problem for the Euler equations to determine
        the intermediate states.
        
        Parameters:
        -----------

        rhol : float
            Density to the left of the discontinuity.

        rhor : float
            Density to the right of the discontinuity

        pl : float
            Pressure to the left of the discontinuity

        pr : float
            Pressure to the right of the discontinuity

        ul : float
            Velocity to the left of the discontinuity

        ur : float
            Velocity to the right of the discontinuity

        Returns:
        --------
    
        (pm,um) : The intermediate states for the Riemann problem.

        """

        gamma = self.gamma

        gamma1 = self.gamma1
        gamma2 = self.gamma2
        gamma3 = self.gamma3
        gamma4 = self.gamma4
        gamma5 = self.gamma5
        gamma6 = self.gamma6
        gamma7 = self.gamma7
        gamma8 = self.gamma8
        
        # sound speed to the left and right based on the Ideal EOS
        cl = sqrt(gamma*pl/rhol)
        cr = sqrt(gamma*pr/rhor)

        # Iteration constants
        tol = self.tol
        niter = self.niter

        # check the initial data
        if ( gamma4*(cl+cr) <= (ur-ul) ):
            raise RuntimeError("Vaccum generated by initial data. Exiting!")

        # compute the guess pressure 'pm'
        qser = 2.0
        cup = 0.25 * (rhol + rhor)*(cl+cr)
        ppv = 0.5 * (pl + pr) + 0.5*(ul - ur)*cup
        pmin = min(pl, pr)
        pmax = max(pl, pr)
        qmax = pmax/pmin

        ppv = max(0.0, ppv)
        
        if ( (qmax <= qser) and (pmin <= ppv) and (ppv <= pmax) ):
            pm = ppv

        elif ( ppv < pmin ):
            
            pq = (pl/pr)**gamma1
            um = ( pq*ul/cl + ur/cr + gamma4*(pq - 1.0) )/( pq/cl + 1.0/cr )
            ptl = 1.0 + gamma7 * (ul - um)/cl
            ptr = 1.0 + gamma7 * (um - ur)/cr
            pm = 0.5*(pl*ptl**gamma3 + pr*ptr**gamma3)
    
        else:

            gel = sqrt( (gamma5/rhol)/(gamma6*pl + ppv) )
            ger = sqrt( (gamma5/rhor)/(gamma6*pr + ppv) )
            pm = (gel*pl + ger*pr - (ur-ul))/(gel + ger)

        # the guessed value is pm
        pstart = pm

        pold = pstart
        udifff = ur-ul

        for i in range(niter):
            fl, fld = self._prefun(pold, rhol, pl, cl)
            fr, frd = self._prefun(pold, rhor, pr, cr)

            p = pold - (fl + fr + udifff)/(fld + frd)
            change = 2.0 * abs((p-pold)/(p+pold))
            if change <= tol:
                break
            pold = p

        if i == niter - 1:
            raise RuntimeError("Divergence in Newton-Raphson Iteration")

        # compute the velocity in the star region 'um'
        um = 0.5 * (ul + ur + fr - fl)
        return p, um

    def sample(self, pm, um, s, rhol, rhor, pl, pr, ul, ur):
        """Sample the solution to the Riemann problem.

        Parameters:
        -----------

        pm, um : float
            Pressure and velocity in the star region as returned
            by `compute_star`

        s : float
            Sampling speed (x/t)

        rhol, rhor : float
            Densities on either side of the discontinuity

        pl, pr : float
            Pressures on either side of the discontinuity

        ul, ur : float
            Velocities on either side of the discontinuity

        Returns:
        --------

        (rho, u, p) : Sampled density, velocity and pressure
        
        """

        gamma = self.gamma

        # sound speeds at the left and right data states
        cl = sqrt(gamma*pl/rhol)
        cr = sqrt(gamma*pr/rhor)

        g1 = self.gamma1 
        g2 = self.gamma2 
        g3 = self.gamma3 
        g4 = self.gamma4
        g5 = self.gamma5
        g6 = self.gamma6
        g7 = self.gamma7
        g8 = self.gamma8

        if s <= um:
            # sampling point lies to the left of the contact discontinuity
            if (pm <= pl):  # left rarefaction

                # speed of the head of the rarefaction
                shl = ul - cl

                if (s <= shl):
                    # sampled point is left state
                    rho = rhol; u = ul; p = pl
                else:
                    cml = cl*(pm/pl)**g1    # eqn (4.54)
                    stl = um - cml          # eqn (4.55)

                    if (s > stl):
                        # sampled point is Star left state. eqn (4.53)
                        rho = rhol*(pm/pl)**(1.0/gamma); u = um; p = pm
                    else:
                    
                        # sampled point is inside left fan
                        u = g5 * (cl + g7*ul + s)
                        c = g5 * (cl + g7*(ul - s))
                        
                        rho = rhol*(c/cl)**g4
                        p = pl * (c/cl)**g3

            else: # pm <= pl
                # left shock
                pml = pm/pl
                sl = ul - cl*sqrt(g2*pml + g1)

                if (s <= sl):
                    # sampled point is left data state
                    rho = rhol; u=ul; p=pl
                else:
                    # sampled point is Star left state
                    rho = rhol*(pml + g6)/(pml*g6 + 1.0)
                    u = um
                    p = pm

        else: # s < um

            # sampling point lise to the right of the contact discontinuity
            if (pm > pr):

                # right shock
                pmr = pm/pr
                sr = ur + cr * sqrt(g2*pmr + g1)

                if (s >= sr):
                    #sampled point is right data state
                    rho = rhor; u = ur; p = pr
                else:
                    # sampled point is star right state
                    rho = rhor*(pmr + g6)/(pmr*g6 + 1.0); u = um; p = pm
            else:
                # right rarefaction
                shr = ur + cr

                if (s >= shr):
                    # sampled point is right state
                    rho = rhor; u = ur; p = pr
                else:
                    cmr = cr*(pm/pr)**g1
                    STR = um + cmr

                    if (s <= STR):
                        # sampled point is star right
                        rho = rhor*(pm/pr)**(1.0/gamma); u = um; p = pm
                    else:
                        # sampled point is inside left fan
                        u = g5*(-cr + g7*ur + s)
                        c = g5*(cr - g7*(ur -s))
                        rho = rhor * (c/cr)**g4
                        p = pr*(c/cr)**g3
                    
        return rho, u, p

    def _prefun(self, p, dk, pk, ck):
        """The pressure function"""
    
        g1 = self.gamma1
        g2 = self.gamma2
        g4 = self.gamma4
        g5 = self.gamma5
        g6 = self.gamma6

        if (p <= pk):
            pratio = p/pk
            f = g4*ck*(pratio**g1 - 1.0)
            fd = (1.0/(dk*ck))*pratio**(-g2)

        else:
            ak = g5/dk
            bk = g6*pk
            qrt = sqrt(ak/(bk+p))
            f = (p-pk)*qrt
            fd = (1.0 - 0.5*(p-pk)/(bk + p))*qrt

        return f, fd

    def godunov_solution(self, wl, wr):
        """Solve the Godunov problem with the given input states."""

        rhol, ul, pl = wl[0], wl[1], wl[2]
        rhor, ur, pr = wr[0], wr[1], wr[2]

        # get the intermediate state
        pm, um = self.compute_star(rhol, rhor, pl, pr, ul, ur)

        # sample the solution
        rho, u, p = self.sample(pm, um, 0, rhol, rhor, pl, pr, ul, ur)

        return rho, u, p

    def set_eulerian_godunov_flux(self, wl, wr, face):
        r"""Set the Eulerian Godunov flux at the interface.

        **Parameters**:

        wl, wr : array
            Primitive variables on either side of the interface.

        flux : array
           The flux array for the interface.

        **Notes**:

        Given the two states :math:`w_l, w_r` of primitive variables
        on either side of the interface, the exact Riemann problem is
        solved to deterine the intermediate state at the interface
        (along the :math:`t` axis). The Eulerian flux is defined as:

        .. math::

            F_{i + \frac{1}{2}} = \begin{matrix} \rho u \\ \rho u^2 +
            p \\ u (E + p) \end{matrix}
                 
        """
        # compute the intermediate state
        rho, u, p = self.godunov_solution(wl, wr)

        # store the flux
        flux = face.flux
        
        flux[0] = rho * u
        flux[1] = flux[0] * u + p
        
        e = p/(rho*self.gamma8); E = rho * (e + 0.5 * u * u)
        flux[2] = u * (E + p)

    def set_lagrangian_godunov_flux(self, wl, wr, face):
        r"""Set the Lagrangian Godunov flux at the interface.

        **Parameters**:

        wl, wr : array
            Primitive variables on either side of the interface.

        flux : array
           The flux array for the interface.

        **Notes**:

        Given the two states :math:`w_l, w_r` of primitive variables
        on either side of the interface, the exact Riemann problem is
        solved to deterine the intermediate state at the interface
        (along the :math:`t` axis). The Eulerian flux is defined as:

        .. math::

            F_{i + \frac{1}{2}} = \begin{matrix} 0 \\ \ p \\ up
            \end{matrix}
                 
        """
        # get the relative velocity at the interface
        uavg = face.compute_roe_average( wl[1], wr[1] )
        uavg = 0
        wl[1] -= uavg; wr[1] -= uavg

        # compute the intermediate state
        rho, u, p = self.godunov_solution(wl, wr)

        # get the absolute velocity
        u += uavg

        # store the face vertex velocity
        face.vertex_vel = u

        # store the flux
        flux = face.flux
        
        flux[0] = 0.0
        flux[1] = p
        flux[2] = u * p

    def compute_waf_states(self, wl, wr, face, U, S):
        """Compute the intermediate states for the WAF method.

        For the WAF method, we need to sample the Riemann problem
        solution for the star-left and star right states. In addition,
        for sonic rarefactions, we also need to sample the Godunov
        solution. In addition, the wave speeds must be set by this
        method which is used in computing the intercell flux.

        Parameters:
        ------------

        wl, wr : array
            Primitive variable states on either side of the cell interface.

        U : array shape = (3, 4) 
            The ouput sampled states are stored in this vector.

        face : hogs.grids.grid1d.Face1D

        S : array shape = (1,3)
            The ouput wave speeds are stored in this vector.

        """

        rhol, ul, pl = wl[0], wl[1], wl[2]
        rhor, ur, pr = wr[0], wr[1], wr[2]

        # get the intermediate state
        pm, um = self.compute_star(rhol, rhor, pl, pr, ul, ur)

        # the contact wave speed is that of the intermediate data
        # state um
        S[1] = um

        # Now determine the nature of the two non-linear waves. This
        # sampling procedure is outlined in section 4.4 of Toro's book

        # sound speeds at the left and right data states
        gamma = self.gamma
        cl = sqrt(gamma*pl/rhol)
        cr = sqrt(gamma*pr/rhor)

        g = self.gamma
        g1 = self.gamma1 
        g2 = self.gamma2 
        g3 = self.gamma3 
        g4 = self.gamma4
        g5 = self.gamma5
        g6 = self.gamma6
        g7 = self.gamma7
        g8 = self.gamma8

        # we begin by checking the left running wave
        
        if pm > pl :   # left shock wave
            pml = pm/pl

            # shock speed. This is the first wave speed 
            sl = ul - cl*sqrt(g2*pml + g1)
            S[0] = sl

            # The data state  is the star-left state
            rho = rhol*(pml + g6)/(pml*g6 + 1.0)
            u = um
            p = pm

        else: # left rarefaction
            
            # check if it is a sonic rarefaction. Compute the speeds
            # of the head and tail and check their signs.
            shl = ul - cl

            cml = cl * (pm/pl)**g1
            stl = um - cml

            # the wave speed is chosen to be that of the head
            S[0] = shl

            if (shl * stl) < 0.0: # is sonic rarefaction

                # the solution is the Godunov's solution
                rho, u, p = self.godunov_solution(wl, wr)

            else: # data state is star left when left wave is a fan
                rho = rhol*(pm/pl)**(1.0/gamma); u = um; p = pm

        w = numpy.array( [rho, u, p] ); q = ideal_gas_w2q(g, w)
        U[:,1] = q

        # now do the analogous thing for the right non-linear wave

        if (pm > pr): # right shock
            
            pmr = pm/pr
            sr = ur + cr * numpy.sqrt(g2*pmr + g1)

            # the wave speed is that of the right shock
            S[2] = sr

            # the data state is star-right
            rho = rhor*(pmr + g6)/(pmr*g6 + 1.0); u = um; p = pm

        else:  # right rarefaction
            
            # compute the speed of the rarefaction head and tail
            shr = ur + cr
            cmr = cr*(pm/pr)**g1
            STR = um + cmr

            # the wave speed is the head speed
            S[2] = shr

            # check if the fan is sonic
            if (shr * STR) < 0.0:

                # the data state is the Godunov's solution
                rho, u, p = self.godunov_solution(wl, wr)

            else: # data state is star right
                rho = rhor*(pm/pr)**(1.0/gamma); u = um; p = pm                
                
        w = numpy.array( [rho, u, p] ); q = ideal_gas_w2q(g, w)
        U[:,2] = q

    def get_flux_jump(self, U, S, i):
        """Get the flux jump across the ith wave

        **Parameters**

        U : array (3,4)
            The state vector for the four states between the 3 waves.

        S : array (1,3)
            Wave speeds for the Riemann solver

        i : int
            The wave number

        """
        Fi = ideal_gas_flux_q( self.gamma, U[:, i] )
        Fip1 = ideal_gas_flux_q( self.gamma, U[:, i + 1] )

        return Fip1 - Fi
